package com.chaosj.skill.zipkin.interceptor

import brave.Tracing
import brave.propagation.Propagation
import brave.propagation.TraceContext
import com.chaosj.skill.zipkin.context.ApplicationContext
import okhttp3.Headers
import okhttp3.Interceptor
import okhttp3.Request
import okhttp3.Response

/**
 * @Description:
 * @author: caoyangjie
 * @date: 2019年12月2019/12/6日 09:32
 * @Copyright: © 赢时胜
 */
class OkhttpTraceInterceptor extends AbstractTraceInterceptor<Headers.Builder> implements Interceptor{
    @Override
    Propagation.Setter<Headers.Builder, String> getSetter() {
        {carrier, key, value -> carrier.add(key,value)}
    }

    @Override
    Response intercept(Chain chain) throws IOException {
        Request request = chain.request()
        TraceContext traceContext = ApplicationContext.getBean(Tracing.class).tracer().currentSpan().context()
        Headers.Builder headers = request.headers().newBuilder()
        headers = injector(headers)
        //rebuild a new request
        request = request.newBuilder().headers(headers.build()).build()
        Response response = chain.proceed(request)
        //将traceId返回去给调用方，如浏览器发起的请求，可在浏览器看到该信息，方便定位
        Headers.Builder responseHeadersBuilder = response.headers()
                .newBuilder()
                .add(TRACE_ID_NAME, traceContext.traceIdString())
        response = response.newBuilder().headers(responseHeadersBuilder.build()).build()

        response
    }
}
